## Load packages
import os
import getpass
import platform
import json
import pandas as pd
import numpy as np
import bipartitepandas as bpd
import pytwoway as tw
import sys

## Set Paths
path = sys.argv[1]
pathr = sys.argv[2]

## Loop through script for each subsample
for sample in ['physicians','lawyers']:
    ## Set parameters
    # FE
    fe_params = tw.fe_params(
        {
            'he': True,
            'ncore': 8,
            'outputfile': pathr+'/'+sample+'/FEestimates_'+sample+'.json'
        }
    )
    # Cleaning
    clean_params = bpd.clean_params(
            {
                'connectedness': 'leave_out_spell',
                'collapse_at_connectedness_measure': True,
                'drop_single_stayers': True,
                'drop_returns': 'returners',
                'copy': False
            }
        )
    
    # Loop through subsamples
    ## Extract data for each subsample
    dtafile = path+'geotwfe03-'+sample+'_movers.dta'
    df = pd.read_stata(dtafile)

    ## Prepare data
    # Convert into BipartitePandas DataFrame
    bdf = bpd.BipartiteDataFrame(
        i=df['i'], j=df['j'], y=df['y'], t=df['t']
    )
    
    # Clean and collapse
    bdf = bdf.clean(clean_params)
    
    ## Initialize and run FE estimator
    fe_estimator = tw.FEEstimator(bdf, fe_params)
    # Fit FE estimator
    fe_estimator.fit()
    
    ## Summarize and output results
    fe_estimator.summary
    results_fe = fe_estimator.summary
    df_fe = pd.DataFrame.from_dict([results_fe]) 
    new_names = {
    'var(y)': 'var_y_',
    'var(eps)_fe': 'var_eps__fe',
    'var(eps)_ho': 'var_eps__ho',
    'var(eps)_he': 'var_eps__he',
    'var(psi)_fe': 'var_psi__fe',
    'var(psi)_ho': 'var_psi__ho',
    'var(psi)_he': 'var_psi__he',
    'cov(psi, alpha)_fe': 'cov_psi__alpha__fe',
    'cov(psi, alpha)_ho': 'cov_psi__alpha__ho',
    'cov(psi, alpha)_he': 'cov_psi__alpha__he'
    }

    with open(pathr+'/'+sample+'/FEestimates_'+sample+'.json') as f:
        data = json.load(f)

    n_workers = int(data['n_workers'])
    df_fe['n_workers'] = n_workers

    df_fe = df_fe.rename(columns=new_names)
    df_fe.to_stata(pathr+'/'+sample+'/geotwfe03-FEestimates_'+sample+'.dta', write_index=False)
